package com.mr.task;

import com.mr.core.*;
import com.mr.job.Job;

import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.Callable;
import java.io.ObjectInputStream;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;

public class ReduceTask<K2, V2, K3, V3> implements Callable<Void> {
    private final List<String> inputFiles;
    private final int partition;
    private final Job<?, ?, K2, V2, K3, V3> job;

    public ReduceTask(List<String> inputFiles, Job<?, ?, K2, V2, K3, V3> job, int partition) {
        this.inputFiles = inputFiles;
        this.job = job;
        this.partition = partition;
    }

    @Override
    public Void call() throws Exception {
        try {
            Reducer<K2, V2, K3, V3> reducer = job.getReducerClass().getDeclaredConstructor().newInstance();
            Context<K3, V3> context = new Context<>();
            
            // 读取所有Map任务的中间结果
            Map<K2, List<V2>> mergedData = new TreeMap<>();
            for (String file : inputFiles) {
                Map<K2, List<V2>> data = readFromFile(file);
                data.forEach((k, v) -> mergedData.merge(k, v, (oldList, newList) -> {
                    oldList.addAll(newList);
                    return oldList;
                }));
            }

            // 执行Reduce（增加异常捕获）
            for (Map.Entry<K2, List<V2>> entry : mergedData.entrySet()) {
                try {
                    reducer.reduce(entry.getKey(), entry.getValue(), context);
                } catch (Exception e) {
                    context.reportError(e);
                }
            }

            // 检查错误状态
            if (context.hasError()) {
                throw new RuntimeException("Reduce task failed", context.getError());
            }

            // 写入最终输出
            String outputFile = String.format("%s/part-%d.out", 
                job.getOutputDir(), partition);
            writeFinalOutput(context.getBuffer(), outputFile);
            return null;
        } catch (Exception e) {
            throw new RuntimeException("Reduce task execution failed", e);
        }
    }

    private Map<K2, List<V2>> readFromFile(String file) {
        try (ObjectInputStream ois = new ObjectInputStream(new FileInputStream(file))) {
            @SuppressWarnings("unchecked") // 确保序列化对象类型安全
            Map<K2, List<V2>> result = (Map<K2, List<V2>>) ois.readObject();
            return result;
        } catch (IOException | ClassNotFoundException e) {
            throw new RuntimeException("Failed to read map output", e);
        }
    }

    private void writeFinalOutput(Map<?, ?> data, String outputFile) throws IOException {
        Path outputPath = Paths.get(outputFile);
        try {
            Files.createDirectories(outputPath.getParent());
            try (PrintWriter writer = new PrintWriter(outputFile, StandardCharsets.UTF_8)) {  
                data.forEach((k, v) -> {
                    // 修改输出格式：当值是列表时取第一个元素
                    if (v instanceof List && !((List<?>)v).isEmpty()) {
                        writer.println(k + "\t" + ((List<?>)v).get(0));
                    } else {
                        writer.println(k + "\t" + v);
                    }
                });
            }
        } catch (IOException e) {
            // 新增详细的错误日志
            System.err.println("Failed to write reduce output: " + e.getMessage());
            throw new RuntimeException("Failed to write output: " + outputFile, e);
        }
    }
}